// ConfidentialityProcessor.java
package com.tidu.demo;

import org.apache.poi.hwpf.HWPFDocument;
import org.apache.poi.hwpf.usermodel.CharacterRun;
import org.apache.poi.hwpf.usermodel.Paragraph;
import org.apache.poi.hwpf.usermodel.Range;
import org.apache.poi.ooxml.POIXMLDocument;
import org.apache.poi.xwpf.usermodel.*;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.Security;
import java.security.spec.KeySpec;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.regex.Pattern;

/**
 * 保密文档处理工具
 * 支持涉密内容识别、替换和加密/解密功能（使用SM4对称加密）
 * 支持DOC和DOCX格式文件
 */
public class ConfidentialityProcessor {

    private static final String SM4_ALGORITHM = "SM4";
    private static final String SM4_TRANSFORMATION = "SM4/ECB/PKCS5Padding";

    // SM4密钥长度
    private static final int SM4_KEY_LENGTH = 128;

    // PBKDF2相关常量
    private static final String PBKDF2_ALGORITHM = "PBKDF2WithHmacSHA256";
    private static final int PBKDF2_ITERATIONS = 10000;
    private static final int KEY_LENGTH = 128;

    // 在静态初始化块中添加BouncyCastle提供者
    static {
        Security.addProvider(new BouncyCastleProvider());
    }

    /**
     * 批量加密文档中的涉密内容
     */
    public static void batchEncryptConfidentialContent(List<String> srcPaths,
                                                       byte[] secretKey,
                                                       ProgressCallback callback) throws Exception {
        for (int i = 0; i < srcPaths.size(); i++) {
            String srcPath = srcPaths.get(i);
            if (callback != null) {
                callback.onProgress(i, srcPaths.size(), "正在加密: " + srcPath);
            }

            try {
                encryptConfidentialContent(srcPath, secretKey);

                if (callback != null) {
                    callback.onSuccess(i, srcPath);
                }
            } catch (Exception e) {
                if (callback != null) {
                    callback.onError(i, srcPath, e);
                }
            }
        }
    }

    /**
     * 批量解密文档中的涉密内容
     */
    public static void batchDecryptConfidentialContent(List<String> srcPaths,
                                                       byte[] secretKey,
                                                       ProgressCallback callback) throws Exception {
        for (int i = 0; i < srcPaths.size(); i++) {
            String srcPath = srcPaths.get(i);
            if (callback != null) {
                callback.onProgress(i, srcPaths.size(), "正在解密: " + srcPath);
            }

            try {
                decryptConfidentialContent(srcPath, secretKey);

                if (callback != null) {
                    callback.onSuccess(i, srcPath);
                }
            } catch (Exception e) {
                if (callback != null) {
                    callback.onError(i, srcPath, e);
                }
            }
        }
    }

    /**
     * 进度回调接口
     */
    public interface ProgressCallback {
        void onProgress(int current, int total, String message);
        void onSuccess(int index, String filePath);
        void onError(int index, String filePath, Exception error);
    }

    /**
     * 生成SM4密钥
     */
    public static byte[] generateSM4Key() throws NoSuchAlgorithmException, NoSuchProviderException {
        KeyGenerator keyGenerator = KeyGenerator.getInstance(SM4_ALGORITHM, "BC");
        keyGenerator.init(SM4_KEY_LENGTH);
        return keyGenerator.generateKey().getEncoded();
    }

    /**
     * 保存密钥到文件
     */
    public static void saveSecretKey(byte[] secretKey, String filename) throws IOException {
        try (FileOutputStream fos = new FileOutputStream(filename)) {
            fos.write(secretKey);
        }
    }

    /**
     * 从文件加载密钥
     */
    public static byte[] loadSecretKey(String filename) throws IOException {
        try (FileInputStream fis = new FileInputStream(filename)) {
            byte[] keyBytes = new byte[SM4_KEY_LENGTH / 8];
            int bytesRead = fis.read(keyBytes);
            if (bytesRead != SM4_KEY_LENGTH / 8) {
                throw new IOException("密钥文件大小不正确");
            }
            return keyBytes;
        }
    }

    /**
     * 从密码生成密钥
     * @param password 用户密码
     * @param salt 盐值（用于增加安全性）
     * @return 生成的密钥字节数组
     */
    public static byte[] deriveKeyFromPassword(String password, byte[] salt) throws Exception {
        KeySpec spec = new PBEKeySpec(password.toCharArray(), salt, PBKDF2_ITERATIONS, KEY_LENGTH);
        SecretKeyFactory factory = SecretKeyFactory.getInstance(PBKDF2_ALGORITHM);
        byte[] keyBytes = factory.generateSecret(spec).getEncoded();
        return keyBytes;
    }

    /**
     * 使用SM4加密内容
     */
    public static byte[] encryptWithSM4(String content, byte[] secretKey) throws Exception {
        Cipher cipher = Cipher.getInstance(SM4_TRANSFORMATION, "BC");
        SecretKeySpec keySpec = new SecretKeySpec(secretKey, SM4_ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, keySpec);
        return cipher.doFinal(content.getBytes("UTF-8"));
    }

    /**
     * 使用SM4解密内容
     */
    public static String decryptWithSM4(byte[] encryptedContent, byte[] secretKey) throws Exception {
        Cipher cipher = Cipher.getInstance(SM4_TRANSFORMATION, "BC");
        SecretKeySpec keySpec = new SecretKeySpec(secretKey, SM4_ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, keySpec);
        byte[] decryptedBytes = cipher.doFinal(encryptedContent);
        return new String(decryptedBytes, "UTF-8");
    }

    /**
     * 加密文档中的涉密内容（直接修改源文件）
     */
    public static void encryptConfidentialContent(String srcPath, byte[] secretKey) throws Exception {
        if (srcPath.toLowerCase().endsWith(".docx")) {
            encryptDocxFile(srcPath, secretKey);
        } else if (srcPath.toLowerCase().endsWith(".doc")) {
            encryptDocFile(srcPath, secretKey);
        } else {
            throw new IllegalArgumentException("不支持的文件格式: " + srcPath);
        }
    }

    /**
     * 解密文档中的涉密内容（直接修改源文件）
     */
    public static void decryptConfidentialContent(String srcPath, byte[] secretKey) throws Exception {
        if (srcPath.toLowerCase().endsWith(".docx")) {
            decryptDocxFile(srcPath, secretKey);
        } else if (srcPath.toLowerCase().endsWith(".doc")) {
            decryptDocFile(srcPath, secretKey);
        } else {
            throw new IllegalArgumentException("不支持的文件格式: " + srcPath);
        }
    }

    /**
     * 加密DOCX文件
     */
    private static void encryptDocxFile(String srcPath, byte[] secretKey) throws Exception {
        XWPFDocument document = null;
        FileOutputStream outStream = null;
        File tempFile = null;

        try {
            // 创建临时文件
            tempFile = File.createTempFile("encrypted_", ".docx");
            tempFile.deleteOnExit();

            // 读取源文档
            document = new XWPFDocument(POIXMLDocument.openPackage(srcPath));

            // 识别涉密内容并直接加密替换
            encryptConfidentialContentInDocument(document, secretKey);

            // 写入加密后的文档到临时文件
            outStream = new FileOutputStream(tempFile);
            document.write(outStream);
            outStream.flush();

            // 关闭所有资源
            if (outStream != null) {
                outStream.close();
                outStream = null;
            }
            if (document != null) {
                document.close();
                document = null;
            }

            // 将临时文件复制回源文件
            try (FileInputStream tempIn = new FileInputStream(tempFile);
                 FileOutputStream srcOut = new FileOutputStream(srcPath)) {
                byte[] buffer = new byte[8192];
                int bytesRead;
                while ((bytesRead = tempIn.read(buffer)) != -1) {
                    srcOut.write(buffer, 0, bytesRead);
                }
            }

        } finally {
            // 清理资源
            if (outStream != null) {
                try {
                    outStream.close();
                } catch (IOException e) {
                    // 忽略
                }
            }
            if (document != null) {
                try {
                    document.close();
                } catch (IOException e) {
                    System.err.println("警告：关闭文档时出错: " + e.getMessage());
                }
            }
            // 删除临时文件
            if (tempFile != null && tempFile.exists()) {
                tempFile.delete();
            }
        }
    }

    /**
     * 解密DOCX文件
     */
    private static void decryptDocxFile(String srcPath, byte[] secretKey) throws Exception {
        XWPFDocument document = null;
        FileOutputStream outStream = null;
        File tempFile = null;

        try {
            // 创建临时文件
            tempFile = File.createTempFile("decrypted_", ".docx");
            tempFile.deleteOnExit();

            // 读取源文档
            document = new XWPFDocument(POIXMLDocument.openPackage(srcPath));

            // 解密文档中的加密内容
            decryptConfidentialContentInDocument(document, secretKey);

            // 写入解密后的文档到临时文件
            outStream = new FileOutputStream(tempFile);
            document.write(outStream);
            outStream.flush();

            // 关闭所有资源
            if (outStream != null) {
                outStream.close();
                outStream = null;
            }
            if (document != null) {
                document.close();
                document = null;
            }

            // 将临时文件复制回源文件
            try (FileInputStream tempIn = new FileInputStream(tempFile);
                 FileOutputStream srcOut = new FileOutputStream(srcPath)) {
                byte[] buffer = new byte[8192];
                int bytesRead;
                while ((bytesRead = tempIn.read(buffer)) != -1) {
                    srcOut.write(buffer, 0, bytesRead);
                }
            }

        } finally {
            // 清理资源
            if (outStream != null) {
                try {
                    outStream.close();
                } catch (IOException e) {
                    // 忽略
                }
            }
            if (document != null) {
                try {
                    document.close();
                } catch (IOException e) {
                    System.err.println("警告：关闭文档时出错: " + e.getMessage());
                }
            }
            // 删除临时文件
            if (tempFile != null && tempFile.exists()) {
                tempFile.delete();
            }
        }
    }

    /**
     * 加密DOC文件
     */
    private static void encryptDocFile(String srcPath, byte[] secretKey) throws Exception {
        HWPFDocument document = null;
        FileOutputStream outStream = null;
        File tempFile = null;

        try {
            // 创建临时文件
            tempFile = File.createTempFile("encrypted_", ".doc");
            tempFile.deleteOnExit();

            // 读取源文档
            document = new HWPFDocument(new FileInputStream(srcPath));

            // 识别涉密内容并直接加密替换
            encryptConfidentialContentInDocDocument(document, secretKey);

            // 写入加密后的文档到临时文件
            outStream = new FileOutputStream(tempFile);
            document.write(outStream);
            outStream.flush();

            // 关闭所有资源
            if (outStream != null) {
                outStream.close();
                outStream = null;
            }
            if (document != null) {
                document.close();
                document = null;
            }

            // 将临时文件复制回源文件
            try (FileInputStream tempIn = new FileInputStream(tempFile);
                 FileOutputStream srcOut = new FileOutputStream(srcPath)) {
                byte[] buffer = new byte[8192];
                int bytesRead;
                while ((bytesRead = tempIn.read(buffer)) != -1) {
                    srcOut.write(buffer, 0, bytesRead);
                }
            }

        } finally {
            // 清理资源
            if (outStream != null) {
                try {
                    outStream.close();
                } catch (IOException e) {
                    // 忽略
                }
            }
            if (document != null) {
                try {
                    document.close();
                } catch (IOException e) {
                    System.err.println("警告：关闭文档时出错: " + e.getMessage());
                }
            }
            // 删除临时文件
            if (tempFile != null && tempFile.exists()) {
                tempFile.delete();
            }
        }
    }

    /**
     * 解密DOC文件
     */
    private static void decryptDocFile(String srcPath, byte[] secretKey) throws Exception {
        HWPFDocument document = null;
        FileOutputStream outStream = null;
        File tempFile = null;

        try {
            // 创建临时文件
            tempFile = File.createTempFile("decrypted_", ".doc");
            tempFile.deleteOnExit();

            // 读取源文档
            document = new HWPFDocument(new FileInputStream(srcPath));

            // 解密文档中的加密内容
            decryptConfidentialContentInDocDocument(document, secretKey);

            // 写入解密后的文档到临时文件
            outStream = new FileOutputStream(tempFile);
            document.write(outStream);
            outStream.flush();

            // 关闭所有资源
            if (outStream != null) {
                outStream.close();
                outStream = null;
            }
            if (document != null) {
                document.close();
                document = null;
            }

            // 将临时文件复制回源文件
            try (FileInputStream tempIn = new FileInputStream(tempFile);
                 FileOutputStream srcOut = new FileOutputStream(srcPath)) {
                byte[] buffer = new byte[8192];
                int bytesRead;
                while ((bytesRead = tempIn.read(buffer)) != -1) {
                    srcOut.write(buffer, 0, bytesRead);
                }
            }

        } finally {
            // 清理资源
            if (outStream != null) {
                try {
                    outStream.close();
                } catch (IOException e) {
                    // 忽略
                }
            }
            if (document != null) {
                try {
                    document.close();
                } catch (IOException e) {
                    System.err.println("警告：关闭文档时出错: " + e.getMessage());
                }
            }
            // 删除临时文件
            if (tempFile != null && tempFile.exists()) {
                tempFile.delete();
            }
        }
    }

    /**
     * 在DOCX文档中加密涉密内容
     */
    private static void encryptConfidentialContentInDocument(XWPFDocument document, byte[] secretKey) throws Exception {
        // 使用配置文件中的关键词构建模式
        Pattern pattern = ConfigManager.buildPattern(ConfigManager.loadCustomKeywords());

        // 加密段落中的涉密内容
        for (XWPFParagraph paragraph : document.getParagraphs()) {
            encryptInParagraph(paragraph, pattern, secretKey);
        }

        // 加密表格中的涉密内容
        for (XWPFTable table : document.getTables()) {
            for (XWPFTableRow row : table.getRows()) {
                for (XWPFTableCell cell : row.getTableCells()) {
                    for (XWPFParagraph paragraph : cell.getParagraphs()) {
                        encryptInParagraph(paragraph, pattern, secretKey);
                    }
                }
            }
        }
    }
    /**
     * 在DOCX段落中加密匹配的内容
     */
    private static void encryptInParagraph(XWPFParagraph paragraph, Pattern pattern, byte[] secretKey) throws Exception {
        // 获取段落的完整文本
        String text = paragraph.getText();
        if (text == null || text.isEmpty()) {
            return;
        }

        // 查找匹配的涉密关键词
        java.util.regex.Matcher matcher = pattern.matcher(text);
        StringBuffer sb = new StringBuffer();
        boolean found = false;

        while (matcher.find()) {
            String match = matcher.group();
            // 检查是否已经是加密占位符
            if (match.startsWith("<E_") && match.endsWith(">")) {
                matcher.appendReplacement(sb, match);
                continue;
            }

            // 对匹配的内容进行加密
            byte[] encryptedBytes = encryptWithSM4(match, secretKey);
            String encryptedData = Base64.getEncoder().encodeToString(encryptedBytes);
            String encryptedPlaceholder = "<E:" + encryptedData + ">";
            matcher.appendReplacement(sb, encryptedPlaceholder);
            found = true;
        }
        matcher.appendTail(sb);

        // 如果有加密内容，更新段落
        if (found) {
            updateParagraphText(paragraph, sb.toString());
        }
    }
    /**
     * 在DOC文档中加密涉密内容
     */
    private static void encryptConfidentialContentInDocDocument(HWPFDocument document, byte[] secretKey) throws Exception {
        // 使用配置文件中的关键词构建模式
        Pattern pattern = ConfigManager.buildPattern(ConfigManager.loadCustomKeywords());

        Range range = document.getRange();

        // 遍历所有段落并加密内容
        for (int i = 0; i < range.numParagraphs(); i++) {
            Paragraph paragraph = range.getParagraph(i);

            // 遍历段落中的所有句子
            for (int j = 0; j < paragraph.numCharacterRuns(); j++) {
                CharacterRun characterRun = paragraph.getCharacterRun(j);
                String text = characterRun.text();

                if (text != null && !text.isEmpty()) {
                    // 查找匹配的涉密关键词
                    java.util.regex.Matcher matcher = pattern.matcher(text);
                    StringBuffer sb = new StringBuffer();
                    boolean found = false;

                    while (matcher.find()) {
                        String match = matcher.group();
                        // 检查是否已经是加密占位符
                        if (match.startsWith("<E_") && match.endsWith(">")) {
                            matcher.appendReplacement(sb, match);
                            continue;
                        }

                        // 对匹配的内容进行加密
                        byte[] encryptedBytes = encryptWithSM4(match, secretKey);
                        String encryptedData = Base64.getEncoder().encodeToString(encryptedBytes);
                        String encryptedPlaceholder = "<E:" + encryptedData + ">";
                        matcher.appendReplacement(sb, encryptedPlaceholder);
                        found = true;
                    }
                    matcher.appendTail(sb);

                    // 如果有加密内容，更新文本
                    if (found) {
                        characterRun.replaceText(text, sb.toString());
                    }
                }
            }
        }
    }
    /**
     * 更新DOCX段落文本
     */
    private static void updateParagraphText(XWPFParagraph paragraph, String newText) {
        // 清空所有runs
        List<XWPFRun> runs = new ArrayList<>(paragraph.getRuns());
        for (XWPFRun run : runs) {
            run.setText("", 0);
        }

        // 如果没有run，创建一个新的
        if (paragraph.getRuns().isEmpty()) {
            paragraph.createRun().setText(newText);
        } else {
            // 在第一个run中设置新文本
            paragraph.getRuns().get(0).setText(newText, 0);
        }
    }

    /**
     * 在DOCX文档中解密加密内容
     */
    private static void decryptConfidentialContentInDocument(XWPFDocument document, byte[] secretKey) throws Exception {
        // 解密段落中的加密内容
        for (XWPFParagraph paragraph : document.getParagraphs()) {
            decryptInParagraph(paragraph, secretKey);
        }

        // 解密表格中的加密内容
        for (XWPFTable table : document.getTables()) {
            for (XWPFTableRow row : table.getRows()) {
                for (XWPFTableCell cell : row.getTableCells()) {
                    for (XWPFParagraph paragraph : cell.getParagraphs()) {
                        decryptInParagraph(paragraph, secretKey);
                    }
                }
            }
        }
    }
    /**
     * 在DOCX段落中解密加密内容
     */
    private static void decryptInParagraph(XWPFParagraph paragraph, byte[] secretKey) throws Exception {
        String text = paragraph.getText();
        if (text == null || text.isEmpty()) {
            return;
        }

        // 查找加密占位符
        Pattern pattern = Pattern.compile("<E:([^>]+)>");
        java.util.regex.Matcher matcher = pattern.matcher(text);
        StringBuffer sb = new StringBuffer();
        boolean found = false;

        while (matcher.find()) {
            String encryptedData = matcher.group(1);
            try {
                // 解密数据
                byte[] encryptedBytes = Base64.getDecoder().decode(encryptedData);
                String decryptedContent = decryptWithSM4(encryptedBytes, secretKey);
                matcher.appendReplacement(sb, decryptedContent);
                found = true;
            } catch (Exception e) {
                // 解密失败，保留原数据
                matcher.appendReplacement(sb, matcher.group());
            }
        }
        matcher.appendTail(sb);

        // 如果有解密内容，更新段落
        if (found) {
            updateParagraphText(paragraph, sb.toString());
        }
    }
    /**
     * 在DOC文档中解密加密内容
     */
    private static void decryptConfidentialContentInDocDocument(HWPFDocument document, byte[] secretKey) throws Exception {
        Range range = document.getRange();
        Pattern pattern = Pattern.compile("<E:([^>]+)>");

        // 遍历所有段落并解密内容
        for (int i = 0; i < range.numParagraphs(); i++) {
            Paragraph paragraph = range.getParagraph(i);

            // 遍历段落中的所有句子
            for (int j = 0; j < paragraph.numCharacterRuns(); j++) {
                CharacterRun characterRun = paragraph.getCharacterRun(j);
                String text = characterRun.text();

                if (text != null && !text.isEmpty()) {
                    // 查找加密占位符
                    java.util.regex.Matcher matcher = pattern.matcher(text);
                    StringBuffer sb = new StringBuffer();
                    boolean found = false;

                    while (matcher.find()) {
                        String encryptedData = matcher.group(1);
                        try {
                            // 解密数据
                            byte[] encryptedBytes = Base64.getDecoder().decode(encryptedData);
                            String decryptedContent = decryptWithSM4(encryptedBytes, secretKey);
                            matcher.appendReplacement(sb, decryptedContent);
                            found = true;
                        } catch (Exception e) {
                            // 解密失败，保留原数据
                            matcher.appendReplacement(sb, matcher.group());
                        }
                    }
                    matcher.appendTail(sb);

                    // 如果有解密内容，更新文本
                    if (found) {
                        characterRun.replaceText(text, sb.toString());
                    }
                }
            }
        }
    }
}
